import pandas as pd
from sklearn.ensemble import AdaBoostRegressor
from sklearn.tree import DecisionTreeClassifier
import util

Xcols = ['eirp',
          'ap_from_ap_0_max_ant_rssi', 'ap_from_ap_0_mean_ant_rssi', 'ap_from_ap_0_sum_ant_rssi',
        # 'ap_from_ap_1_max_ant_rssi', 'ap_from_ap_1_mean_ant_rssi', 'ap_from_ap_1_sum_ant_rssi',
          'sta_from_ap_0_max_ant_rssi', 'sta_from_ap_0_mean_ant_rssi', 'sta_from_ap_0_sum_ant_rssi',
          'sta_from_ap_1_max_ant_rssi', 'sta_from_ap_1_mean_ant_rssi', 'sta_from_ap_1_sum_ant_rssi',
          'sta_to_ap_0_max_ant_rssi', 'sta_to_ap_0_mean_ant_rssi', 'sta_to_ap_0_sum_ant_rssi',
          'sta_to_ap_1_max_ant_rssi', 'sta_to_ap_1_mean_ant_rssi', 'sta_to_ap_1_sum_ant_rssi']

classModel = util.loadModel('mcsClass.pkl')
regModel = util.loadModel('mcsReg.pkl')

def predTestCsv(filename):
    testData = pd.read_csv(filename)
    testX = testData[Xcols]
    testY = classModel.predict(testX)
    testY = testY.tolist()
    print(testY)
    print('classModel result')
    testX = testX.values.tolist()
    for i in range(len(testY)):
        b11 = testY[i]
        if not b11:
            # 不是11的调回归模型预测
            newMcsValue = regModel.predict([testX[i]])
            testY[i] = int(round(newMcsValue[0]))
        else:
            testY[i] = 11
    print(testY)
    print('regModel result')
    testData['mcs'] = pd.Series(testY)
    testData.to_csv(filename)

predTestCsv('test_set_2_2ap_ReduceRssi_pred.csv')
predTestCsv('test_set_2_3ap_ReduceRssi_pred.csv')